/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __SFT_ESTIMATE_WORKER_HH__
#define __SFT_ESTIMATE_WORKER_HH__

#include <boost/optional.hpp>
#include <boost/interprocess/detail/atomic.hpp>
#include <cstdlib>
#include "example.hh"
#include "learningrate.hh"
#include "loss.hh"
#include "regressor.hh"
#include "threadqueue.hh"

namespace flassol
{
  class SftEstimateWorker
    {
      private:
        CSMCQueue&      in;
        CSMCQueue&      out;
        Regressor&      regressor;
        uint64_t        thread_no;
        uint64_t        num_threads;
        LearningRate&   eta;
        Loss&           loss;
        float           lambda;
        bool            adaptive;
        bool            test_only;

        double          t;
        unsigned int    num_classes;

        static float
        square (float x)
          {
            return x * x;
          }

        static inline void
        spin_barrier (volatile boost::uint32_t* x)
          {
            using boost::interprocess::detail::atomic_dec32;
            using boost::interprocess::detail::atomic_read32;

            atomic_dec32 (x);

            while (atomic_read32 (x) > 0)
              {
              }
          }

        void
        do_estimate (CSMCExample&          ex,
                     EstimateResultVector& result,
                     unsigned int          level)
          {
            EstimateResultMatrix& matrix = 
              (level == 0) ? (*ex.initials) : (*ex.partials);
            EstimateResultVector& partial = matrix[thread_no];

            if (level > 1)
              {
                for (unsigned int k = 0; k + 1 < num_classes; ++k)
                  {
                    partial[k].reset ();
                    partial[k].skip = result[k].skip;
                  }
              }

            regressor.get_partial (ex.f.get (),
                                   thread_no,
                                   num_threads,
                                   partial);

            for (unsigned int k = 0; k + 1 < num_classes; ++k)
              {
                result[k] = partial[k];
              }

            if (num_threads > 1)
              {
                spin_barrier (ex.estimationcount + level);

                for (uint64_t i = 0; i < num_threads; ++i)
                  {
                    if (i != thread_no)
                      {
                        for (unsigned int k = 0; k + 1 < num_classes; ++k)
                          {
                            result[k] += matrix[i][k];
                          }
                      }
                  }
              }
          }

        void
        do_adaptive_norm (CSMCExample&        ex,
                          UpdateInfoVector&   info,
                          unsigned int        level)
          {
            regressor.get_adaptive_norm_partial (ex.f.get (),
                                                 thread_no,
                                                 num_threads,
                                                 info);

            if (num_threads > 1)
              {
                for (unsigned int k = 0; k + 1 < num_classes; ++k)
                  {
                    (*ex.norm_partials)[thread_no][k] = info[k].adaptivexnorm;
                  }

                spin_barrier (ex.adaptivecount + level);

                for (uint64_t i = 0; i < num_threads; ++i)
                  {
                    if (i != thread_no)
                      {
                        for (unsigned int k = 0; k + 1 < num_classes; ++k)
                          {
                            info[k].adaptivexnorm += 
                              (*ex.norm_partials)[thread_no][k];
                          }
                      }
                  }
              }
          }

        void
        do_update (CSMCExample&                 ex,
                   const UpdateInfoVector&      info,
                   unsigned int                 level)
          {
            regressor.update_partial (ex.f.get (),
                                      info,
                                      thread_no,
                                      num_threads);

            if (num_threads > 1)
              {
                spin_barrier (ex.updatecount + level);
              }
          }

        void
        do_filter_tree (CSMCExample& ex)
          {
            EstimateResultVector result (num_classes - 1);

            do_estimate (ex, result, 0);

            if (ex.cost && ! test_only)
              {
                SparseVector cost (*ex.cost);

                if (result[0].xnorm > 0)
                  {
                    ++t;

                    float thiseta = eta.eta (t);
                    unsigned int winner[num_classes];

                    for (unsigned int l = 0; l < num_classes; ++l)
                      {
                        winner[l] = l;
                      }

                    for (unsigned int l = num_classes, level = 0;
                         l > 1 && ! cost.empty ();
                         l = ((l >> 1) + (l % 2)), ++level)
                      {
                        UpdateInfoVector info (num_classes - 1);

                        if (level > 0)
                          {
                            do_estimate (ex, result, level);
                          }

                        for (unsigned int n = 0; 2 * n + 1 < l; ++n)
                          {
                            unsigned int left = winner[2 * n];
                            unsigned int right = winner[2 * n + 1];

                            float leftp = 
                              (left == 0) ? 0 : result[left - 1].p;
                            float rightp = 
                              (right == 0) ? 0 : result[right - 1].p;
                            float thisp = leftp - rightp;

                            float thisimp = 
                              fabsf (cost[left] - cost[right]);

                            float thislabel = 
                              (cost[left] < cost[right])
                                ? 1 : -1;

                            float gsq = 
                              square (  loss.dldp (thisp, thislabel)
                                      * thisimp);
                            
                            if (left > 0)
                              {
                                info[left - 1].gsq = gsq;
                              }

                            if (right > 0)
                              {
                                info[right - 1].gsq = gsq;
                              }
                          }

                        if (adaptive)
                          {
                            do_adaptive_norm (ex, info, level);
                          }

                        for (unsigned int n = 0; 
                             2 * n + 1 < l && ! cost.empty (); 
                             ++n)
                          { 
                            unsigned int left = winner[2 * n];
                            unsigned int right = winner[2 * n + 1];

                            float leftp =
                              (left == 0) ? 0 : result[left - 1].p;
                            float rightp =
                              (right == 0) ? 0 : result[right - 1].p;

                            float thisimp =
                              fabsf (cost[left] - cost[right]);

                            if (thisimp > 0 &&
                                info[std::max (left, right) - 1].gsq > 0)
                              { 
                                float leftwdotx =
                                  (left == 0) ? 0 : result[left - 1].wdotx;
                                float rightwdotx =
                                  (right == 0) ? 0 : result[right - 1].wdotx;
                                float leftadotb =
                                  (left == 0) ? 0 : result[left - 1].adotb;
                                float rightadotb =
                                  (right == 0) ? 0 : result[right - 1].adotb;
                                float leftanormplusbnorm =
                                  (left == 0)
                                    ? 0 : result[left - 1].anormplusbnorm;
                                float rightanormplusbnorm =
                                  (right == 0)
                                    ? 0 : result[right - 1].anormplusbnorm;
                                float leftadaptivexnorm =
                                  (left == 0)
                                    ? 0 : info[left - 1].adaptivexnorm;
                                float rightadaptivexnorm =
                                  (right == 0)
                                    ? 0 : info[right - 1].adaptivexnorm;

                                float thisp = leftp - rightp;
                                float thiswdotx = leftwdotx - rightwdotx;
                                float thisadotb = leftadotb - rightadotb;
                                float thisanormplusbnorm =
                                  leftanormplusbnorm + rightanormplusbnorm;
                                float thisadaptivexnorm =
                                  leftadaptivexnorm + rightadaptivexnorm;

                                float thislabel =
                                  (cost[left] < cost[right])
                                    ? 1 : -1;

                                std::pair<float, float> update =
                                  loss.invariant_update
                                    (thisp,
                                     thiswdotx,
                                     thisadotb,
                                     thislabel,
                                     thisimp * thiseta,
                                     thisanormplusbnorm,
                                     2.0f * 
                                       result[std::max (left, 
                                                        right) - 1].xnorm,
                                     adaptive,
                                     thisadaptivexnorm,
                                     lambda);

                                if (left > 0)
                                  { 
                                    info[left - 1].sh = update.first;
                                    info[left - 1].decay = update.second;
                                  }

                                if (right > 0)
                                  { 
                                    info[right - 1].sh = -update.first;
                                    info[right - 1].decay = update.second;
                                  }
                              }

                            if (leftp > rightp)
                              { 
                                winner[n] = left;
                                cost[right] = 0;
                                if (right > 0)
                                  {
                                    result[right - 1].skip = true;
                                  }
                              }
                            else
                              { 
                                winner[n] = right;
                                cost[left] = 0;
                                if (left > 0)
                                  {
                                    result[left - 1].skip = true;
                                  }
                              }
                          }

                        if (l % 2)
                          { 
                            winner[(l - 1) / 2] = winner[l - 1];
                          }

                        do_update (ex, info, level);
                      }
                  }
              }
          }

      public:
        SftEstimateWorker (CSMCQueue&           _in,
                           CSMCQueue&           _out,
                           Regressor&           _regressor,
                           uint64_t             _thread_no,
                           uint64_t             _num_threads,
                           LearningRate&        _eta,
                           Loss&                _loss,
                           float                _lambda,
                           bool                 _adaptive,
                           bool                 _test_only)
          : in (_in),
            out (_out),
            regressor (_regressor),
            thread_no (_thread_no),
            num_threads (_num_threads),
            eta (_eta),
            loss (_loss),
            lambda (_lambda),
            adaptive (_adaptive),
            test_only (_test_only),
            t (0),
            num_classes (regressor.get_header ().num_classes)
          {
          }

        void operator() ()
          {
            std::vector<boost::optional<CSMCExample> >* rvec;

            out.add_writer ();

            while (in.pop (rvec))
              {
                using boost::interprocess::detail::atomic_dec32;

                for (unsigned int n = 0; n < rvec->size (); ++n)
                  {
                    if ((*rvec)[n])
                      {
                        do_filter_tree (*(*rvec)[n]);
                      }
                  }

                if (! rvec->back ())
                  {
                    if (thread_no == 0)
                      {
                        out.push (rvec);
                      }
                  }
                else if (num_threads == 1 || 
                         atomic_dec32 (rvec->back ()->usecount) == 1)
                  {
                    out.push (rvec);
                  }
              }

            out.del_writer ();
          }
    };
}

#endif // __SFT_ESTIMATE_WORKER_HH__
